import os
import sys
import warnings
import numpy as np
import pandas as pd
from scipy.signal import welch, find_peaks
import antropy as ant  # 保留 antropy 引用，以便未来恢复使用
from tqdm import tqdm

# 如果你有自定义 pymagic 路径，加入 sys.path
dirs = ['/database/home/duansizhang/pywave']
for d in dirs:
    if d not in sys.path:
        sys.path.insert(0, d)

from pymagic.tools.ecgdllinterface import DataSource
import matplotlib.pyplot as plt

# ———— 配置区域 ————
raw_file = '/database/private/mgcdb/raw/0a0b34b8220644848209d7a199afd6fc_1699056000000.raw'
out_dir   = '/database/home/duansizhang/hrv_predict/result'
os.makedirs(out_dir, exist_ok=True)
warnings.filterwarnings('ignore')
# ————————————————————

# 1. 读取加速度数据
hdata = DataSource(raw_file, 3)
acc = hdata.GetAccData()  # Nx3 array
N = acc.shape[0]

# 2. 参数设置
fs = 10  # 实际采样率 10 Hz
window_sec = 300  # 窗口长度 300 秒
window_len = window_sec * fs  # 75,000 样本
nperseg = fs * 300  # 10 秒段用于 Welch 频谱估计
num_windows = N // window_len

# 3. 特征提取（已移除最慢的非线性熵特征，以加速）
features = []
for wi in tqdm(range(num_windows), desc="Processing windows"):
    start = wi * window_len
    end = start + window_len
    seg = acc[start:end, :]

    # 时域特征
    sma = np.sum(np.abs(seg)) / window_len
    std_mean = np.std(seg, axis=0).mean()
    peak_total = sum(len(find_peaks(seg[:, i])[0]) for i in range(3))

    # 频域特征（使用较小 nperseg 提速）
    f, Pxx = welch(seg, fs=fs, axis=0, nperseg=nperseg)
    fft_power = np.sum(Pxx, axis=0).mean()
    spec_centroid = (f[:, None] * Pxx).sum(axis=0).mean() / fft_power
    spec_flux = np.sqrt(np.sum(np.diff(Pxx, axis=0)**2, axis=0)).mean()

    # 汇总
    features.append({
        'win_start_s': start / fs,
        'sma': sma,
        'std_mean': std_mean,
        'peak_total': peak_total,
        'fft_power': fft_power,
        'spec_centroid': spec_centroid,
        'spec_flux': spec_flux
    })

# 4. 保存 CSV
df = pd.DataFrame(features)
out_csv = os.path.join(out_dir, 'acc_features.csv')
df.to_csv(out_csv, index=False)
print(f"Saved acceleration features to: {out_csv}")

# 5. 可选：绘制 SMA 趋势图
plt.figure(figsize=(10, 4))
plt.plot(df['win_start_s'], df['sma'], marker='o')
plt.title('Windowed Signal Magnitude Area (SMA)')
plt.xlabel('Window start (s)')
plt.ylabel('SMA')
plt.tight_layout()
plt.savefig(os.path.join(out_dir, 'acc_sma_over_time.png'))
plt.close()
